from modelling import *


class Modeller:
    def __init__(self, modelling_func, **config):
        assert 'model_type' in config and 'batch_size' in config, 'shuffle' in config
        self.model_type = config['model_type']
        train_dataset, test_dataset, self.network, self.loss_fn, self.validate_func\
            = modelling_func(config['dataset_config'] if 'dataset_config' in config else {})
        self.train_loader = paddle.io.DataLoader(train_dataset, batch_size=config['batch_size'],
                                                 shuffle=config['shuffle'] if 'shuffle' in config else True)
        self.test_loader = paddle.io.DataLoader(test_dataset, batch_size=config['batch_size'],
                                                shuffle=config['shuffle'] if 'shuffle' in config else True)
        self.validation_config = config

    def train(self, epochs=5, optimizer=None, validate_every=100):
        assert self.train_loader and self.test_loader and self.network and self.loss_fn
        if optimizer is None:
            optimizer = paddle.optimizer.Adam(learning_rate=0.001, parameters=self.network.parameters())

        for epoch in range(epochs):
            for batch_id, (x_data, y_data) in enumerate(self.train_loader()):
                predicts = self.network(x_data)
                acc = paddle.metric.accuracy(predicts, y_data)
                loss = self.loss_fn(predicts, y_data)
                loss.backward()
                if batch_id % validate_every == 0:
                    print("epoch: {}, batch_id: {}, loss: {}, acc: {}".format(epoch, batch_id, loss.numpy(),
                                                                              acc.numpy()))
                    self.validate_func(self.network(next(self.test_loader())[1][0]), self.validation_config)
                optimizer.step()
                optimizer.clear_grad()

    def train_gan(self, epochs=5, optimizer=None, validate_every=100):
        assert self.train_loader and self.test_loader and self.network and self.loss_fn
        if optimizer is None:
            optimizer = paddle.optimizer.Adam(learning_rate=0.001, parameters=self.network.parameters())

        for epoch in range(epochs):
            for batch_id, (x_data, y_data) in enumerate(self.train_loader()):
                predicts = self.network(x_data)
                acc = paddle.metric.accuracy(predicts, y_data)
                loss = self.loss_fn(predicts, y_data)
                loss.backward()
                if batch_id % validate_every == 0:
                    print("epoch: {}, batch_id: {}, loss: {}, acc: {}".format(epoch, batch_id, loss.numpy(),
                                                                              acc.numpy()))
                    self.validate_func(self.network(next(self.test_loader())[1][0]), self.validation_config)
                optimizer.step()
                optimizer.clear_grad()

            for batch_id, (x_data, y_data) in enumerate(self.train_loader()):
                predicts = self.network(x_data)
                acc = paddle.metric.accuracy(predicts, y_data)
                loss = self.loss_fn(predicts, y_data)
                loss.backward()
                if batch_id % validate_every == 0:
                    print("epoch: {}, batch_id: {}, loss: {}, acc: {}".format(epoch, batch_id, loss.numpy(),
                                                                              acc.numpy()))
                    self.validate_func(self.network(next(self.test_loader())[1][0]), self.validation_config)
                optimizer.step()
                optimizer.clear_grad()


if __name__ == '__main__':
    task_type = 'audio2params'

    if task_type == 'audio2params':
        dataset_config = {
            'train_sample_dir': '',
            'test_sample_dir': '',
            'train_label_dir': '',
            'test_label_dir': '',
            'plugin_path': '',
            'preset_file_format': 'fxp',
            'compress_rate': 100.,
        }

        validation_config = {
            'source_sample_rate': 44100.,
        }

        modeller = Modeller(audio2params_modelling,
                            model_type='normal',
                            batch_size=16,
                            dataset_config=dataset_config,
                            validation_config=validation_config)
        modeller.train(epochs=5)

    elif task_type == 'piece_gan':
        dataset_config = {
            'train_sample_dir': '',
            'test_sample_dir': '',
            'train_label_dir': '',
            'test_label_dir': '',
            'plugin_path': '',
            'preset_file_format': 'fxp',
            'compress_rate': 100.,
        }

        validation_config = {
            'source_sample_rate': 44100.,
        }

        modeller = Modeller(audio2params_modelling,
                            model_type='normal',
                            batch_size=16,
                            dataset_config=dataset_config,
                            validation_config=validation_config)
        modeller.train_gan(epochs=5)
